#pragma once

#include <iostream>
using namespace std;

enum Colour
{
    BLACK,
    RED
};

template <class T>
struct RBTreeNode
{
    RBTreeNode<T>* _left;
    RBTreeNode<T>* _right;
    RBTreeNode<T>* _parent;
    Colour _col;
    T _data;

    RBTreeNode(const T& data)
        : _left(nullptr)
        , _right(nullptr)
        , _parent(nullptr)
        , _col(RED)
        , _data(data)
    {}
};

template <class T, class Ref, class Ptr>
struct RBTreeIterator
{
    typedef RBTreeNode<T> Node;
    typedef RBTreeIterator<T, Ref, Ptr> Self;

    Node* _node;

    RBTreeIterator(Node* node = nullptr)
        :_node(node)
    {}

    Ref operator*()
    {
        return _node->_data;
    }

    Ptr operator->()
    {
        return &(_node->_data);
    }

    Self& operator++()
    {
        if(_node->_right)   //右子树不为空，找右子树中序第一个节点
        {
            Node* subLeft = _node->_right;
            while(subLeft->_left)
            {
                subLeft = subLeft->_left;
            }
            _node = subLeft;
        }
        else    //右子树为空，说明当前子树已经访问完了，沿着到根的路径向上，找孩子是父亲的左的那个父亲节点
        {
            Node* cur = _node;
            Node* parent = cur->_parent;
            while (parent && parent->_right == cur)
            {
                cur = parent;
                parent = parent->_parent;
            }
            _node = parent;
        } 
        return *this;
    }

    bool operator!=(const Self& it) const
    {
        return this->_node != it._node;
    }

    bool operator==(const Self& it) const
    {
        return this->_node == it._node;
    }

};

//T->K: set
//T->pair<const k, v>: map
template <class K, class T, class KeyOfT>
class RBTree
{
    typedef RBTreeNode<T> Node;
public:
    typedef RBTreeIterator<T, T&, T*> iterator;
    typedef RBTreeIterator<const T, const T&, const T*> const_iterator;
public:
    iterator begin()
    {
        Node* left = _root;
        while(left && left->_left)
        {
            left = left->_left;
        }
        return left;    //单参数的构造函数支持隐式类型转换
        // return iterator(left);
    }

    iterator end()
    {
        // return iterator(nullptr);
        return nullptr;
    }
private:
    Node* _root;
public:
    RBTree(): _root(nullptr)
    {}
    pair<iterator, bool> Insert(const T& data)
    {
        if(_root == nullptr)
        {
            _root = new Node(data);
            _root->_col = BLACK;
            return make_pair(iterator(_root), true);
        }

        Node* parent = nullptr;
        Node* cur = _root;
        KeyOfT kot;
        
        while(cur != nullptr)
        {
            if(kot(cur->_data) < kot(data))
            {
                parent = cur;
                cur = cur->_right;
            }
            else if(kot(cur->_data) > kot(data))
            {
                parent = cur;
                cur = cur->_left;
            }
            else
            {
                return make_pair(iterator(cur), false);
            }
        }
        cur = new Node(data);
        Node* newnode = cur;
        // cur->_col = RED;
        if(kot(parent->_data) > kot(data))
        {
            parent->_left = cur;
            cur->_parent = parent;
        }
        else
        {
            parent->_right = cur;
            cur->_parent = parent;
        }

        //balance control
        while(parent != nullptr && parent->_col == RED)
        {
            Node* grandfather = parent->_parent;
            if(parent == grandfather->_left)
            {
                Node* uncle = grandfather->_right;
                //情况一：叔叔存在且为红，进行变色处理，并继续向上处理
                if(uncle && uncle->_col == RED)
                {
                    parent->_col = uncle->_col = BLACK;
                    grandfather->_col = RED;

                    cur = grandfather;
                    parent = cur->_parent;
                }
                else    //uncle不存在 or 存在且为黑
                {
                    if(cur == parent->_left)    //情况二：单旋 + 变色
                    {
                        RotateR(grandfather);
                        parent->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    else    //情况三：双旋 + 变色
                    {
                        RotateL(parent);
                        RotateR(grandfather);
                        grandfather->_col = RED;
                        cur->_col = BLACK;
                    }
                    break;
                }
            }
            else
            {
                Node* uncle = grandfather->_left;
                if(uncle && uncle->_col == RED)
                {
                    parent->_col = BLACK;
                    uncle->_col = BLACK;
                    grandfather->_col = RED;

                    cur = grandfather;
                    parent = cur->_parent;
                }
                else
                {
                    if(cur == parent->_right)
                    {
                        RotateL(grandfather);
                        parent->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    else
                    {
                        RotateR(parent);
                        RotateL(grandfather);
                        cur->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    break;
                }
            }
        }
        _root->_col = BLACK;
        return make_pair(iterator(newnode), true);
    }

    void RotateL(Node* parent)
    {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
        parent->_right = subRL;
        if(subRL != nullptr)
        {
            subRL->_parent = parent;
        }
        Node* ppNode = parent->_parent;
        subR->_left = parent;
        parent->_parent = subR;

        if(parent == _root)
        {
            _root = subR;
            subR->_parent = nullptr;
        }
        else
        {
            if(ppNode->_left == parent)
            {
                ppNode->_left = subR;
                subR->_parent = ppNode;
            }
            else
            {
                ppNode->_right = subR;
                subR->_parent = ppNode;
            }
        }
        
    }

    void RotateR(Node* parent)
    {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
        parent->_left = subLR;
        if(subLR != nullptr)
        {
            subLR->_parent = parent;
        }
        Node* ppNode = parent->_parent;
        subL->_right = parent;
        parent->_parent = subL;

        if(parent == _root)
        {
            _root = subL;
            subL->_parent = nullptr;
        }
        else
        {
            if(ppNode->_left == parent)
            {
                ppNode->_left = subL;
            }
            else
            {
                ppNode->_right = subL;
            }
            subL->_parent = ppNode;
        }
    }

    bool IsBalance()
    {
        if(_root == nullptr)
        {
            return true;
        }
        if(_root->_col == RED)
        {
            cerr << "error: root is RED!" << endl;
            return false;
        }

        //算出最左路径的黑色节点作为一个基准值
        Node* cur = _root;
        int benchmark = 0;
        while(cur)
        {
            if(cur->_col == BLACK)
            {
                benchmark++;
            }
            cur = cur->_left;
        }

        return CheckRED_RED(_root)
            && CheckBlackNum(_root, 0, benchmark);
    }

    void InOrder()
    {
        _InOrder(_root);
    }
  private:
    void _InOrder(Node* root)
    {
        KeyOfT kot;
        if(root == nullptr)
        {
            return;
        }
        _InOrder(root->_left);
        cout << kot(root->_data) << endl;
        _InOrder(root->_right);
    }
    bool CheckRED_RED(Node* cur)
    {
        if(cur == nullptr)
        {
            return true;
        }

        if(cur->_col == RED && cur->_parent->_col == RED)
        {
            cerr << "error: 连续红节点！" << endl;
            return false; 
        }

        return CheckRED_RED(cur->_left)
            && CheckRED_RED(cur->_right);
    }
    bool CheckBlackNum(Node* cur, int blackNum, int benchmark)
    {
        if(cur == nullptr)
        {
            if(blackNum != benchmark)
            {
                cerr << "error: 每条路径上的黑色节点不完全相同！" << endl;
                return false;
            }
            return true;
        }
        if(cur->_col == BLACK)
        {
            blackNum++;
        }
        return CheckBlackNum(cur->_left, blackNum, benchmark)
            && CheckBlackNum(cur->_right, blackNum, benchmark);
    }
};
